import cv2, time
import numpy as np

"""
使用光流估计方法提取运动信息
"""
MAX_POINT_COUNT = 100

lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))

feature_params = dict(maxCorners=500,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)


def GetBackGround():
    cap = cv2.VideoCapture('sources/FixedPosVideo.avi')
    tracks = []
    last_gray = None
    track_len = 10
    detect_interval = 5
    frame_idx = 0

    while (1):
        ret, frame = cap.read()
        if ret == True:
            gray = cv2.cvtColor(frame, code=cv2.COLOR_BGR2GRAY)
            vis = frame.copy()

            if len(tracks) > 0:  # 检测到角点后进行光流跟踪
                img0, img1 = last_gray, gray

                p0 = np.float32([tr[-1] for tr in tracks]).reshape(-1, 1, 2)

                # 前一帧的角点和当前帧的图像作为输入来得到角点在当前帧的位置
                # prevImg, nextImg, prevPts, nextPts,
                p1, st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params)

                # 当前帧跟踪到的角点及图像和前一帧的图像作为输入来找到前一帧的角点位置
                p0r, st, err = cv2.calcOpticalFlowPyrLK(img1, img0, p1, None, **lk_params)

                d = abs(p0 - p0r).reshape(-1, 2).max(-1)  # 得到角点回溯与前一帧实际角点的位置变化关系

                good = d < 1  # 判断d内的值是否小于1，大于1跟踪被认为是错误的跟踪点

                new_tracks = []
                for tr, (x, y), good_flag in zip(tracks, p1.reshape(-1, 2), good):  # 将跟踪正确的点列入成功跟踪点
                    if not good_flag:
                        continue
                    tr.append((x, y))
                    if len(tr) > track_len:
                        del tr[0]
                    new_tracks.append(tr)
                    cv2.circle(vis, (x, y), 2, (0, 255, 0), -1)
                tracks = new_tracks
                # 以上一帧角点为初始点，当前帧跟踪到的点为终点划线
                cv2.polylines(vis, [np.int32(tr) for tr in tracks], False, (0, 255, 0))

            if frame_idx % detect_interval == 0:  # 每隔5帧检测一次特征点
                mask = np.zeros_like(gray)  # 初始化和适配大小相同的图像
                mask[:] = 255  # 将mask赋值255也就是算全部图像的角点
                for x, y in [np.int32(tr[-1]) for tr in tracks]:  # 跟踪的角点画圆
                    cv2.circle(mask, (x, y), 5, 0, -1)
                p = cv2.goodFeaturesToTrack(gray, mask=mask, **feature_params)  # 像素级别角点检测

                if p is not None:
                    for x, y in np.float32(p).reshape(-1, 2):
                        tracks.append([(x, y)])  # 将检测到的角点放在待跟踪的序列中

            frame_idx += 1
            last_gray = gray
            cv2.imshow('lk_track', vis)
            cv2.imshow('mask', mask)

        k = cv2.waitKey(30) & 0xff
        if k == 27:
            break

    cap.release()
    cv2.destroyAllWindows()


def getArea(image, x, y, w, h):
    """
    获取图片的一个区域的图形 ,返回面积(像素值)
    image 是一个二值图，信息点是白色
    """
    zone = image[y:y + h, x:x + w]
    area = 0

    area = np.sum(zone == 255)  # 等价于后面的循环
    # for row in zone:
    #     for col in row:
    #         if col == 255:
    #             area += 1
    # print("area: ", area)
    return area


if __name__ == '__main__':
    GetBackGround()
    # help(cv2.goodFeaturesToTrack)
    a = np.array([2, 1, 4, 3, 5, 6, 9, 8, 7]).reshape(-1, 3)
    print(a)
    b = a.max(2)
    print(b)
